package consensus

import (
	"log"
	"sync"
	"time"
)

// EpochManager manages epoch execution similar to Turritopsis
type EpochManager struct {
	// Configuration
	epochs        int
	r             int  // ADKR execution interval
	n             int  // Number of nodes
	f             int  // Number of faulty nodes
	pid           int  // Process ID
	batchSize     int  // Batch size for transactions
	
	// Components
	acs           ACSInterface
	adkr          ADKRInterface
	consensus     *EbbFlowConsensus
	
	// Metrics
	metrics       *Metrics
	epochResults  []*EpochResult
	
	// Synchronization
	mu            sync.RWMutex
}

// Metrics tracks performance metrics
type Metrics struct {
	TotalEpochs     int
	CompletedEpochs int
	FailedEpochs    int
	TotalTPS        float64
	AverageLatency  time.Duration
}

// EpochResult stores the result of an epoch
type EpochResult struct {
	EpochID      int
	StartTime    time.Time
	EndTime      time.Time
	Success      bool
	TPS          float64
	Latency      time.Duration
	Error        error
}

// ACSInterface defines the interface for ACS
type ACSInterface interface {
	Execute(batchSize int) error
}

// ADKRInterface defines the interface for ADKR
type ADKRInterface interface {
	Execute() error
}

// NewEpochManager creates a new epoch manager
func NewEpochManager(epochs, r, n, f, pid, batchSize int) *EpochManager {
	return &EpochManager{
		epochs:      epochs,
		r:           r,
		n:           n,
		f:           f,
		pid:         pid,
		batchSize:   batchSize,
		metrics:     &Metrics{},
		epochResults: make([]*EpochResult, 0),
	}
}

// SetACS sets the ACS component
func (em *EpochManager) SetACS(acs ACSInterface) {
	em.acs = acs
}

// SetADKR sets the ADKR component
func (em *EpochManager) SetADKR(adkr ADKRInterface) {
	em.adkr = adkr
}

// SetConsensus sets the consensus component
func (em *EpochManager) SetConsensus(consensus *EbbFlowConsensus) {
	em.consensus = consensus
}

// Run executes the epoch management
func (em *EpochManager) Run() error {
	log.Printf("Starting epoch management: %d epochs, R=%d", em.epochs, em.r)
	
	for epoch := 0; epoch < em.epochs; epoch++ {
		startTime := time.Now()
		
		// Execute ACS
		if err := em.executeACS(epoch); err != nil {
			log.Printf("ACS execution failed for epoch %d: %v", epoch, err)
			em.updateMetrics(epoch, startTime, time.Now(), false, 0, 0, err)
			continue
		}
		
		// Execute ADKR if needed
		if epoch%em.r == 0 && epoch > 0 {
			if err := em.executeADKR(epoch); err != nil {
				log.Printf("ADKR execution failed for epoch %d: %v", epoch, err)
			}
		}
		
		endTime := time.Now()
		latency := endTime.Sub(startTime)
		tps := float64(em.batchSize) / latency.Seconds()
		
		em.updateMetrics(epoch, startTime, endTime, true, tps, latency, nil)
		
		log.Printf("Epoch %d completed: TPS=%.2f, Latency=%v", epoch, tps, latency)
		
		// Small delay between epochs
		time.Sleep(100 * time.Millisecond)
	}
	
	em.printFinalStatistics()
	return nil
}

// executeACS executes the ACS protocol
func (em *EpochManager) executeACS(epoch int) error {
	if em.acs == nil {
		// Simulate ACS execution
		time.Sleep(50 * time.Millisecond)
		return nil
	}
	
	return em.acs.Execute(em.batchSize)
}

// executeADKR executes the ADKR protocol
func (em *EpochManager) executeADKR(epoch int) error {
	if em.adkr == nil {
		// Simulate ADKR execution
		time.Sleep(100 * time.Millisecond)
		return nil
	}
	
	return em.adkr.Execute()
}

// updateMetrics updates the performance metrics
func (em *EpochManager) updateMetrics(epoch int, startTime, endTime time.Time, success bool, tps float64, latency time.Duration, err error) {
	em.mu.Lock()
	defer em.mu.Unlock()
	
	result := &EpochResult{
		EpochID:   epoch,
		StartTime: startTime,
		EndTime:   endTime,
		Success:   success,
		TPS:       tps,
		Latency:   latency,
		Error:     err,
	}
	
	em.epochResults = append(em.epochResults, result)
	em.metrics.TotalEpochs++
	
	if success {
		em.metrics.CompletedEpochs++
		em.metrics.TotalTPS += tps
		em.metrics.AverageLatency += latency
	} else {
		em.metrics.FailedEpochs++
	}
}

// printFinalStatistics prints the final statistics
func (em *EpochManager) printFinalStatistics() {
	em.mu.RLock()
	defer em.mu.RUnlock()
	
	if em.metrics.CompletedEpochs > 0 {
		avgTPS := em.metrics.TotalTPS / float64(em.metrics.CompletedEpochs)
		avgLatency := em.metrics.AverageLatency / time.Duration(em.metrics.CompletedEpochs)
		
		log.Printf("=== Epoch Management Statistics ===")
		log.Printf("Total Epochs: %d", em.metrics.TotalEpochs)
		log.Printf("Completed: %d", em.metrics.CompletedEpochs)
		log.Printf("Failed: %d", em.metrics.FailedEpochs)
		log.Printf("Average TPS: %.2f", avgTPS)
		log.Printf("Average Latency: %v", avgLatency)
	}
}